import numpy as np
from utils import uniform_samples

def ransac_polyfit(x_sub, y_sub, degree, num_iter=100, threshold=1.0):
    # from sklearn.linear_model import RANSACRegressor, LinearRegression
    # model = RANSACRegressor(estimator=LinearRegression(), max_trials=num_iter, residual_threshold=threshold)
    
    # # 需要将数据转变为二维数组
    # x_sub_reshaped = x_sub[:, np.newaxis]
    # try:
    #     model.fit(np.vander(x_sub, degree + 1), y_sub)
    #     coefficients = model.estimator_.coef_
    #     coefficients[0] = model.estimator_.intercept_
    #     best_fit = coefficients[::-1]  # reverse to match np.polyval direction
    # except ValueError as e:
    #     print(f"RANSAC fit failed due to: {e}")
    #     return None

    # return best_fit
    best_fit = None
    max_inliers = 0

    # 检查是否有足够的点用于采样
    if len(x_sub) < degree + 1:
        print("Not enough data points to perform sampling.")
        return None

    for _ in range(num_iter):
        sample_indices = np.random.choice(len(x_sub), degree + 1, replace=False)
        x_sample = x_sub[sample_indices]
        y_sample = y_sub[sample_indices]

        try:
            A = np.vander(x_sample, degree + 1, increasing=True)
            coefficients = np.linalg.solve(np.dot(A.T, A), np.dot(A.T, y_sample))
            
            y_estimated = np.polyval(coefficients[::-1], x_sub)
            residuals = np.abs(y_sub - y_estimated)

            inliers = residuals < threshold
            num_inliers = np.sum(inliers)

            if num_inliers > max_inliers:
                max_inliers = num_inliers
                best_fit = coefficients
        except np.linalg.LinAlgError:
            continue

    return best_fit

def run_curve_fitting(x_samples, y_samples_noisy, intervals=None, num_experiments=1000, max_degree=4, use_ransac=False):
    """
    最小二乘拟合，返回拟合结果。
    """
    if intervals is None:
        intervals = [1]  

    fit_results = {}

    for num_intervals in intervals:
        sub_interval_results = []
        x_samples = np.array(x_samples)
        interval_length = (x_samples.max() - x_samples.min()) / num_intervals

        print(f"Number of Intervals: {num_intervals}, Interval Length: {interval_length}")

        for i in range(num_intervals):
            # 计算每个区间的边界
            start = x_samples.min() + i * interval_length
            end = start + interval_length

            x_sub = x_samples[(x_samples >= start) & (x_samples < end)]
            y_sub = y_samples_noisy[(x_samples >= start) & (x_samples < end)]

            print(f"Interval {i}: x_sub = {x_sub}, y_sub = {y_sub}")

            if len(x_sub) == 0 or len(y_sub) == 0:
                print(f"Empty subarray for interval {i}, skipping...")
                continue

            x_experiment = uniform_samples(start, end, num_experiments)


            for degree in range(1, max_degree + 1):
                try:
                    A = np.vander(x_sub, degree + 1, increasing=True)
                    A_T_A = np.dot(A.T, A)
                    A_T_b = np.dot(A.T, y_sub)

                    coefficients = np.linalg.solve(A_T_A, A_T_b)
                    y_fit = np.polyval(coefficients[::-1], x_experiment)

                    # RANSAC 拟合
                    if use_ransac:
                        coefficients_ransac = ransac_polyfit(x_sub, y_sub, degree)
                        y_fit_ransac = np.polyval(coefficients_ransac[::-1], x_experiment)
                        y_fit = y_fit_ransac

                    sub_interval_results.append({
                        'degree': degree,
                        'y_fit': y_fit
                    })
                except np.linalg.LinAlgError as e:
                    print(f"Linear algebra error for degree {degree} in interval {i}: {e}")

            fit_results[num_intervals] = sub_interval_results

    return fit_results
